# coding:gbk
import time
from datetime import datetime
from collections import defaultdict

# ==================== 常量定义层 ====================
class Constants:
    """集中管理所有策略常量，避免魔法数字"""
    
    # 交易操作类型 (根据QMT API文档定义)
    OP_TYPE_BUY = 23      # 买入操作码
    OP_TYPE_SELL = 24     # 卖出操作码
    
    # 交易所代码
    EXCHANGE_ID_STOCK = 1101  # A股市场代码
    
    # 价格类型
    PRICE_TYPE_LIMIT = 11     # 限价单
    
    # 订单状态码 (需与券商系统保持一致)
    STATUS_UNREPORTED = 48         # 未报
    STATUS_WAITING = 49            # 待报
    STATUS_REPORTED = 50           # 已报
    STATUS_REPORTED_CANCEL = 51    # 已报待撤
    STATUS_PART_FILLED_CANCEL = 52 # 部成待撤
    STATUS_CANCELLED = 54          # 已撤
    STATUS_FILLED = 56             # 已成
    STATUS_INVALID = 57            # 废单

    # 状态分组（用于逻辑判断）
    ORDER_INITIAL_STATUSES = {
        STATUS_UNREPORTED, STATUS_WAITING, STATUS_REPORTED,
        STATUS_REPORTED_CANCEL, STATUS_PART_FILLED_CANCEL
    }
    ORDER_FINAL_STATUSES = {
        STATUS_FILLED, STATUS_PART_FILLED_CANCEL, 
        STATUS_CANCELLED, STATUS_INVALID
    }
    ORDER_SUCCESS_STATUSES = {STATUS_FILLED, STATUS_PART_FILLED_CANCEL}

    # 交易方向（用于订单匹配）
    TRADE_DIRECTION_BUY = 1
    TRADE_DIRECTION_SELL = 2

    # 风控参数
    MAX_DRAWDOWN = -0.05          # 单日最大回撤5%
    MAX_TRADES_PER_DAY = 20       # 单标的日最大交易次数
    MIN_EXEC_INTERVAL = 1         # 最小执行间隔(秒)
    ORDER_TIMEOUT = 10            # 订单超时时间(秒) 
    CANCEL_RETRY_TIMES = 3        # 撤单重试次数
    CANCEL_RETRY_INTERVAL = 2     # 撤单重试间隔(秒)

# ==================== 核心组件层 ====================
class DataCache:
    """
    高频数据缓存器
    功能：减少API调用频率，提升性能
    """
    def __init__(self):
        self.positions = None    # 缓存持仓数据
        self.orders = None       # 缓存订单数据
        self.last_update = 0     # 最后更新时间戳
        
    def get_positions(self, force_update=False):
        """获取持仓数据（带缓存）"""
        if force_update or time.time() - self.last_update > 5:
            self._refresh_all()
        return self.positions
    
    def get_orders(self, force_update=False):
        """获取订单数据（带缓存）"""
        if force_update or time.time() - self.last_update > 5:
            self._refresh_all()
        return self.orders
    
    def _refresh_all(self):
        """强制刷新所有数据"""
        self.positions = get_trade_detail_data(g.accountID, 'stock', 'POSITION')
        self.orders = get_trade_detail_data(g.accountID, 'stock', 'ORDER')
        self.last_update = time.time()

class OrderMonitor:
    """
    订单全生命周期管理器
    改进点：增加超时撤单和异常订单检测
    """
    def __init__(self):
        # 结构：{order_id: (stock_code, timestamp, volume, direction)}
        self.pending_orders = {}  
    
    def add_order(self, order_id, stock_code, volume, direction):
        """注册新订单并记录方向"""
        self.pending_orders[order_id] = (
            stock_code, 
            time.time(), 
            volume,
            direction
        )
    
    def check_orders(self, ContextInfo):
        """
        订单状态检查主逻辑
        处理超时订单和异常状态
        """
        current_orders = get_trade_detail_data(g.accountID, 'stock', 'ORDER')
        
        for order_id, (stock_code, ts, exp_vol, direction) in list(self.pending_orders.items()):
            order = self._find_order(current_orders, order_id)
            
            # 订单查询失败处理
            if not order:
                if time.time() - ts > Constants.ORDER_TIMEOUT:
                    self._handle_timeout(order_id, stock_code, ContextInfo)
                continue
                
            # 终态订单处理
            if order.m_nOrderStatus in Constants.ORDER_FINAL_STATUSES:
                self._process_final_status(order, stock_code, exp_vol)
                self.pending_orders.pop(order_id)
    
    def has_pending_order(self, stock_code, direction):
        """检查同方向未完成订单（防重复下单）"""
        for order_id, (code, _, _, dir_) in self.pending_orders.items():
            if code == stock_code and dir_ == direction:
                return True
        return False
    
    def _find_order(self, orders, target_id):
        """精确查找订单"""
        if not isinstance(orders, list): return None
        for o in orders:
            if hasattr(o, 'm_strOrderSysID') and o.m_strOrderSysID == target_id:
                return o
        return None
    
    def _process_final_status(self, order, stock_code, exp_vol):
        """处理终态订单（成功/失败）"""
        if order.m_nOrderStatus in Constants.ORDER_SUCCESS_STATUSES:
            handle_trade_success(stock_code, order.m_dTradedPrice)
        else:
            handle_trade_failure(stock_code)
    
    def _handle_timeout(self, order_id, stock_code, ContextInfo):
        """
        超时订单处理（核心改进）
        执行智能撤单流程
        """
        print(f"! 订单超时尝试撤单: {order_id}")
        if smart_cancel(order_id, stock_code, ContextInfo):
            reset_and_unlock(stock_code, "超时撤单成功")
        else:
            print(f"! 撤单失败保持锁定: {order_id}")
            # 记录异常供监控系统捕获
            g.trading_states[stock_code]['last_error'] = '撤单超时'

class RiskControl:
    """多维度风控引擎"""
    def __init__(self):
        self.trade_counts = defaultdict(int)  # 单日交易计数
        self.last_trade_time = {}             # 最后交易时间
    
    def check_interval(self, stock_code):
        """检查最小交易间隔"""
        now = time.time()
        last_time = self.last_trade_time.get(stock_code, 0)
        return now - last_time >= Constants.MIN_EXEC_INTERVAL
    
    def record_trade(self, stock_code):
        """记录交易事件"""
        self.trade_counts[stock_code] += 1
        self.last_trade_time[stock_code] = time.time()

# ==================== 全局状态层 ====================
class GlobalState:
    """策略运行时状态容器"""
    def __init__(self):
        self.base_prices = {}      # 各标的基准价
        self.trading_states = {}   # 交易状态机
        self.accountID = ''        # 交易账户
        self.last_checked_date = None  # 最后检查日期
        self.components = {       # 核心组件实例
            'monitor': OrderMonitor(),
            'risk': RiskControl(),
            'cache': DataCache()
        }
        self.last_exec_time = 0   # 上次执行时间戳

g = GlobalState()

# ==================== 策略逻辑层 ====================
def init(ContextInfo):
    """策略初始化入口"""
    ContextInfo.account_id = '39135479'  # 替换为真实账户
    
    # 参数配置（根据标的特性调整）
    ContextInfo.stock_params = {
        '510900.SH': {  # H股ETF
            'initial_volume': 5000,     # 初始仓位(股)
            'initial_price': 1.087,   # 初始基准价
            'grid_gap': 0.4,          # 网格间距(%)
            'trade_buffer': 500,      # 交易缓冲股数
            'pullback_trigger_pct': 0.01,  # 触发观察模式的涨跌幅
            'pullback_entry_pct': 0.001    # 重新激活的回撤幅度
        },
        '513060.SH': {  # 恒生医疗ETF
            'initial_volume': 10000,
            'initial_price': 0.541,
            'grid_gap': 0.4,
            'trade_buffer': 1000,
            'pullback_trigger_pct': 0.01,
            'pullback_entry_pct': 0.001
        }
    }
    
    # 初始化标的列表
    ContextInfo.stock_list = list(ContextInfo.stock_params.keys())
    ContextInfo.trade_start_time = datetime.strptime("09:30:00", "%H:%M:%S").time()
    ContextInfo.trade_end_time = datetime.strptime("15:00:00", "%H:%M:%S").time()

    # 初始化全局状态
    g.accountID = ContextInfo.account_id
    g.last_checked_date = datetime.now().date()
    
    for stock_code, params in ContextInfo.stock_params.items():
        g.base_prices[stock_code] = params['initial_price']
        g.trading_states[stock_code] = {
            'mode': 'TRADING',        # 交易模式（TRADING/OBSERVING）
            'order_locked': False,    # 订单锁定状态
            'locked_order_id': None,  # 锁定订单ID
            'locked_timestamp': 0,    # 锁定时间戳
            'last_peak': params['initial_price'],  # 观察期最高价
            'last_trough': params['initial_price'], # 观察期最低价
            'last_error': None        # 最后错误信息
        }

    # 设置标的并刷新缓存
    ContextInfo.set_universe(ContextInfo.stock_list)
    g.components['cache']._refresh_all()
    
    print("策略初始化完成：超时撤单增强版 v8.2")
    print(f"监控标的: {ContextInfo.stock_list}")

def handlebar(ContextInfo):
    """主事件循环（每个tick触发）"""
    now = datetime.now()
    
    # ---- 1. 交易时间过滤 ----
    if not (ContextInfo.trade_start_time <= now.time() <= ContextInfo.trade_end_time):
        return
    
    # ---- 2. 执行频率控制 ----
    current_time = time.time()
    if current_time - g.last_exec_time < Constants.MIN_EXEC_INTERVAL:
        return
    g.last_exec_time = current_time
    
    # ---- 3. 交易日切换检测 ----
    today = now.date()
    if today != g.last_checked_date:
        print(f"新交易日 {today}，重置状态...")
        for stock_code in ContextInfo.stock_list:
            g.trading_states[stock_code].update({
                'mode': 'TRADING',
                'last_error': None
            })
            g.components['risk'].trade_counts[stock_code] = 0
        g.last_checked_date = today
    
    # ---- 4. 执行订单状态检查 ----
    g.components['monitor'].check_orders(ContextInfo)
    
    # ---- 5. 异常订单检测（每分钟）----
    if int(time.time()) % 60 == 0:
        detect_abnormal_orders(ContextInfo)
    
    # ---- 6. 获取缓存数据 ----
    ticks = ContextInfo.get_full_tick(ContextInfo.stock_list)
    positions = g.components['cache'].get_positions()
    position_map = {pos.m_strInstrumentID: pos.m_nVolume 
                   for pos in positions if hasattr(pos, 'm_strInstrumentID')}
    
    # ---- 7. 标的轮询 ----
    for stock_code in ContextInfo.stock_list:
        if stock_code not in ticks or 'lastPrice' not in ticks[stock_code]:
            continue
            
        tick_data = ticks[stock_code]
        state = g.trading_states[stock_code]
        
        # ---- 8. 风控检查 ----
        if not (g.components['risk'].check_interval(stock_code) and
                not state.get('last_error')):
            continue
            
        # ---- 9. 模式判断 ----
        if state['order_locked']:
            continue  # 锁定状态由OrderMonitor处理
        elif state['mode'] == 'OBSERVING':
            check_pullback_and_reactivate(ContextInfo, stock_code, tick_data['lastPrice'])
        elif state['mode'] == 'TRADING':
            current_position = position_map.get(stock_code, 0)
            run_grid_logic(ContextInfo, stock_code, tick_data, current_position)

# ==================== 交易逻辑函数 ====================
def run_grid_logic(ContextInfo, stock_code, tick_data, current_position):
    """
    执行网格交易逻辑
    包含防重复下单和撤单保护机制
    """
    params = ContextInfo.stock_params[stock_code]
    
    # 1. 计算动态目标仓位
    base_volume = current_position if current_position > 0 else params['initial_volume']
    price_diff = (g.base_prices[stock_code] - tick_data['lastPrice']) 
    dynamic_gap = params['grid_gap'] * (1 + abs(price_diff)/g.base_prices[stock_code])
    target_position = base_volume + price_diff * 1000 * dynamic_gap
    
    # 2. 获取买卖盘口（防数组越界）
    ask_prices = tick_data.get('askPrice', [0]*5)
    bid_prices = tick_data.get('bidPrice', [0]*5)
    buy_price = ask_prices[1] if len(ask_prices) > 1 else ask_prices[0]
    sell_price = bid_prices[1] if len(bid_prices) > 1 else bid_prices[0]
    
    # 3. 交易决策
    order_amount, order_price, op_type = 0, 0, 0
    if target_position > current_position + params['trade_buffer']:
        order_amount = min(
            int((target_position - current_position)/100)*100,  # 整手数
            tick_data.get('askVol', [0]*5)[1]  # 不超过卖一量
        )
        order_price = buy_price
        op_type = Constants.OP_TYPE_BUY
    elif target_position < current_position - params['trade_buffer']:
        order_amount = min(
            int((current_position - target_position)/100)*100,
            current_position  # 不超过当前持仓
        )
        order_price = sell_price
        op_type = Constants.OP_TYPE_SELL
    
    # 4. 执行交易（带防护机制）
    if order_amount > 0 and order_price > 0:
        direction = Constants.TRADE_DIRECTION_BUY if op_type == Constants.OP_TYPE_BUY else Constants.TRADE_DIRECTION_SELL
        
        # 关键防护1：检查同方向未完成订单
        if g.components['monitor'].has_pending_order(stock_code, direction):
            print(f"防重复检查: {stock_code} 存在未完成{'买入' if op_type == Constants.OP_TYPE_BUY else '卖出'}订单，跳过")
            return
            
        print(f"交易信号: {stock_code} {'买' if op_type == Constants.OP_TYPE_BUY else '卖'} {order_amount}股 @ {order_price}")
        
        # 发送订单
        passorder(
            op_type, Constants.EXCHANGE_ID_STOCK, g.accountID, 
            stock_code, Constants.PRICE_TYPE_LIMIT, order_price, 
            order_amount, '', 1, ContextInfo
        )
        
        # 关键防护2：记录交易并锁定状态
        temp_order_id = f"TEMP_{stock_code}_{time.time()}"
        g.components['monitor'].add_order(
            temp_order_id, stock_code, order_amount, direction
        )
        g.components['risk'].record_trade(stock_code)
        set_lock(stock_code, temp_order_id)

# ==================== 撤单相关函数 ====================
def smart_cancel(order_id, stock_code, ContextInfo):
    """
    智能撤单流程
    返回bool表示是否成功
    """
    for i in range(Constants.CANCEL_RETRY_TIMES):
        try:
            # 1. 检查订单最新状态
            order = get_value_by_order_id(order_id, g.accountID, 'stock', 'ORDER')
            if not order or order.m_nOrderStatus in Constants.ORDER_FINAL_STATUSES:
                print(f"订单 {order_id} 已终态，无需撤单")
                return True
            
            # 2. 执行撤单
            if ContextInfo.cancel(order_id, g.accountID, 'stock'):
                print(f"第{i+1}次撤单成功: {order_id}")
                return True
                
        except Exception as e:
            print(f"第{i+1}次撤单异常: {e}")
        
        time.sleep(Constants.CANCEL_RETRY_INTERVAL)
    
    print(f"! 撤单彻底失败: {order_id}")
    return False

def detect_abnormal_orders(ContextInfo):
    """
    定时检测异常挂单
    处理长时间锁定的订单
    """
    for stock_code, state in g.trading_states.items():
        if (state['order_locked'] and 
            time.time() - state['locked_timestamp'] > Constants.ORDER_TIMEOUT * 3):
            print(f"! 检测到异常挂单: {stock_code} 锁定超过{Constants.ORDER_TIMEOUT * 3}秒")
            if smart_cancel(state['locked_order_id'], stock_code, ContextInfo):
                reset_and_unlock(stock_code, "异常订单已清理")

# ==================== 状态管理函数 ====================
def set_lock(stock_code, order_id):
    """锁定订单状态"""
    state = g.trading_states[stock_code]
    state.update({
        'order_locked': True,
        'locked_order_id': order_id,
        'locked_timestamp': time.time(),
        'last_error': None
    })
    print(f"锁定 {stock_code} 订单: {order_id}")

def reset_and_unlock(stock_code, reason=""):
    """
    安全解锁函数
    包含状态最终检查
    """
    state = g.trading_states[stock_code]
    if state['locked_order_id']:
        # 最终检查订单是否已成交
        order = get_value_by_order_id(state['locked_order_id'], g.accountID, 'stock', 'ORDER')
        if order and order.m_nOrderStatus == Constants.STATUS_FILLED:
            handle_trade_success(stock_code, order.m_dTradedPrice)
            return
    
    # 正常解锁流程
    state.update({
        'order_locked': False,
        'locked_order_id': None,
        'locked_timestamp': 0
    })
    print(f"解锁 {stock_code} | 原因: {reason}")

def handle_trade_success(stock_code, traded_price):
    """成功交易处理"""
    state = g.trading_states[stock_code]
    g.base_prices[stock_code] = traded_price  # 更新基准价
    state.update({
        'last_peak': traded_price,
        'last_trough': traded_price,
        'mode': 'OBSERVING',
        'last_error': None
    })
    print(f"标的 {stock_code} 基准价更新为 {traded_price}，进入观察模式")
    reset_and_unlock(stock_code, "交易成功")

def handle_trade_failure(stock_code):
    """失败交易处理"""
    state = g.trading_states[stock_code]
    state.update({
        'mode': 'TRADING',
        'last_error': None
    })
    print(f"标的 {stock_code} 交易失败，返回交易模式")
    reset_and_unlock(stock_code, "交易失败")

# ==================== 系统回调函数 ====================
def deal_callback(ContextInfo, dealInfo):
    """成交回报处理"""
    print(f"成交回报: {dealInfo.m_strInstrumentID} "
          f"{dealInfo.m_dPrice}@{dealInfo.m_nVolume}股")
    # 强制刷新缓存
    g.components['cache']._refresh_all()